package org.dcarew.stripjar;

import org.apache.bcel.classfile.ClassParser;
import org.apache.bcel.classfile.Code;
import org.apache.bcel.classfile.Constant;
import org.apache.bcel.classfile.ConstantClass;
import org.apache.bcel.classfile.ConstantPool;
import org.apache.bcel.classfile.JavaClass;
import org.apache.bcel.classfile.Method;
import org.apache.bcel.generic.ConstantPoolGen;
import org.apache.bcel.generic.INVOKESTATIC;
import org.apache.bcel.generic.Instruction;
import org.apache.bcel.generic.InstructionList;
import org.apache.bcel.generic.LDC;
import org.apache.bcel.generic.Type;
import org.apache.commons.compress.archivers.zip.ZipArchiveEntry;
import org.apache.commons.compress.archivers.zip.ZipFile;

import java.io.IOException;
import java.util.HashSet;
import java.util.Set;
import java.util.zip.ZipException;

class ClassEntriesCollector {
  private ZipFile zipFile;

  private Set<String> toVisit = new HashSet<String>();
  private Set<String> visited = new HashSet<String>();
  private Set<String> contains = new HashSet<String>();

  public ClassEntriesCollector(ZipFile zipFile, String entryPointClass) throws ZipException,
      IOException {
    this.zipFile = zipFile;

    addToVisit(entryPointClass);
  }

  public void addToWhitelist(String className) {
    addToVisit(className);
  }

  public void collect() throws IOException {
    while (!toVisit.isEmpty()) {
      String className = toVisit.iterator().next();

      toVisit.remove(className);

      visit(className);
    }
  }

  public boolean contains(ZipArchiveEntry entry) {
    return contains.contains(entry.getName());
  }

  void addToVisit(String className) {
    if (!visited.contains(className)) {
      toVisit.add(className);
    }
  }

  void visit(String className) throws IOException {
    visited.add(className);

    ZipArchiveEntry entry = zipFile.getEntry(Utils.classNameToPathName(className));

    if (entry != null) {
      contains.add(entry.getName());

      ClassParser parser = new ClassParser(zipFile.getInputStream(entry), entry.getName());
      JavaClass javaClass = parser.parse();

      // traverse the constant pool, looking for class entries
      ConstantPool cp = javaClass.getConstantPool();

      for (Constant constant : cp.getConstantPool()) {
        if (constant instanceof ConstantClass) {
          ConstantClass classConstant = (ConstantClass) constant;

          String referencedClass = classConstant.getBytes(cp);

          if (!referencedClass.startsWith("[")) {
            addToVisit(referencedClass.replace('/', '.'));
          }
        }
      }

      ConstantPoolGen cpGen = new ConstantPoolGen(cp);

      // check for Class.forName()
      for (Method method : javaClass.getMethods()) {
        Code code = method.getCode();

        if (code != null) {
          InstructionList instrList = new InstructionList(code.getCode());

          Instruction prevInst = null;

          for (Instruction instr : instrList.getInstructions()) {
            if (instr instanceof INVOKESTATIC) {
              INVOKESTATIC invokeInstr = (INVOKESTATIC) instr;

              if (isClassForName(cpGen, invokeInstr)) {
                Type[] argTypes = invokeInstr.getArgumentTypes(cpGen);

                if (argTypes.length == 1 && Type.STRING.equals(argTypes[0])) {
                  if (isStringInst(cpGen, prevInst)) {
                    LDC ldc = (LDC) prevInst;
                    String str = (String) ldc.getValue(cpGen);

                    addToVisit(str);

                    System.out.println("  found Class.forName(): " + str);
                  }
                }
              }
            }

            prevInst = instr;
          }
        }
      }
    }
  }

  private boolean isClassForName(ConstantPoolGen cp, INVOKESTATIC invokeInstr) {
    // static method, java.lang.Class.forName(java.lang.String)java.lang.Class;

    String methodName = invokeInstr.getName(cp);

    if (!methodName.equals("forName")) {
      return false;
    }

    Type callType = invokeInstr.getReferenceType(cp);
    Type retType = invokeInstr.getReturnType(cp);

    if (!Type.CLASS.equals(callType)) {
      return false;
    }

    if (!Type.CLASS.equals(retType)) {
      return false;
    }

    return true;
  }

  private boolean isStringInst(ConstantPoolGen cp, Instruction inst) {
    if (!(inst instanceof LDC)) {
      return false;
    }

    LDC ldc = (LDC) inst;

    return Type.STRING.equals(ldc.getType(cp));
  }

}
